import os
import sys
sys.path.append(".")

import pyarma as pa

from unifuncs import *

def Euler1(u0, h, A, g):
    
    c = [0]
    return u0 + h * pv.phiv(h * A, [g(u0) + A * u0])

def StrehmelStrong2(u0, h, A, g):

    c = [0, 1]
    Un = []
    Gn = []

    Un.append(u0)    
    Gn.append(g(Un[0]))

    Un.append(u0 + c[1] * h * pv.phiv(c[1] * h * A, [Gn[0] + A * u0]))
    Gn.append(g(Un[1]))

    return u0 + h * pv.phiv(h * A, [Gn[0] + A * u0, - Gn[0] + Gn[1]])

def StrehmelWeak2(u0, h, A, g):

    c = [0, 1/2]
    Un = []
    Gn = []

    Un.append(u0)    
    Gn.append(g(Un[0]))

    Un.append(u0 + c[1] * h * pv.phiv(c[1] * h * A, [Gn[0] + A * u0]))
    Gn.append(g(Un[1]))

    return u0 + h * pv.phiv(h * A, [Gn[1] + A * u0])

def Cox4(u0, h, A, g):

    c = [0, 1/2, 1/2, 1]

    Un = []
    Gn = []

    Un.append(u0)
    Gn.append(g(Un[0]))

    Un.append(u0 + c[1] * h * pv.phiv(c[1] * h * A, [Gn[0] + A * u0]))
    Gn.append(g(Un[1]))

    Un.append(u0 + c[2] * h * pv.phiv(c[2] * h * A, [Gn[1] + A * u0]))
    Gn.append(g(Un[2]))

    Un.append(u0 + c[3] * h * pv.phiv(c[2] * h * A, [(1/2) * (pa.expmat(c[2] * h * A) * Gn[0] - Gn[0]) + Gn[2] + A * u0]))
    Gn.append(g(Un[3]))

    return u0 + h * pv.phiv(h * A, [Gn[0] + A * u0, -3 * Gn[0] + 2 * Gn[1] + 2 * Gn[2] - Gn[3], 4 * Gn[0] - 4 * Gn[1] - 4 * Gn[2] + 4 * Gn[3]])

def Krogstad4(u0, h, A, g):

    c = [0, 1/2, 1/2, 1]
    Un = []
    Gn = []

    Un.append(u0)
    Gn.append(g(Un[0]))

    Un.append(u0 + c[1] * h * pv.phiv(c[1] * h * A, [Gn[0] + A * u0]))
    Gn.append(g(Un[1]))

    Un.append(u0 + c[2] * h * pv.phiv(c[2] * h * A, [Gn[0] + A * u0, -2 * Gn[0] + 2 * Gn[1]]))
    Gn.append(g(Un[2]))

    Un.append(u0 + c[3] * h * pv.phiv(c[3] * h * A, [Gn[0] + A * u0, -2 * Gn[0] + 2 * Gn[2]]))
    Gn.append(g(Un[3]))

    return u0 + h * pv.phiv(h * A, [Gn[0] + A * u0, -3 * Gn[0] + 2 * Gn[1] + 2 * Gn[2] - Gn[3], 4 * Gn[0] - 4 * Gn[1] - 4 * Gn[2] + 4 * Gn[3]])

def Strehmel4(u0, h, A, g):

    c = [0, 1/2, 1/2, 1]
    Un = []
    Gn = []

    Un.append(u0)
    Gn.append(g(Un[0]))

    Un.append(u0 + c[1] * h * pv.phiv(c[1] * h * A, [Gn[0] + A * u0]))
    Gn.append(g(Un[1]))

    Un.append(u0 + c[2] * h * pv.phiv(c[2] * h * A, [Gn[0] + A * u0, -Gn[0] + Gn[1]]))
    Gn.append(g(Un[2]))

    Un.append(u0 + c[3] * h * pv.phiv(c[3] * h * A, [Gn[0] + A * u0, -2 * Gn[0] - 2 * Gn[1] + 4 * Gn[2]]))
    Gn.append(g(Un[3]))

    return u0 + h * pv.phiv(h * A, [Gn[0] + A * u0, -3 * Gn[0] + 4 * Gn[2] - Gn[3], 4 * Gn[0] - 8 * Gn[2] + 4 * Gn[3]])

def Hochbruck5(u0, h, A, g):

    c = [0, 1/2, 1/2, 1, 1/2]
    Un = []
    Gn = []

    Un.append(u0)
    Gn.append(g(Un[0]))

    Un.append(u0 + c[1] * h * pv.phiv(c[1] * h * A, [Gn[0] + A * u0]))
    Gn.append(g(Un[1]))

    Un.append(u0 + c[2] * h * pv.phiv(c[2] * h * A, [Gn[0] + A * u0, -2 * Gn[0] + 2 * Gn[1]]))
    Gn.append(g(Un[2]))

    Un.append(u0 + c[3] * h * pv.phiv(c[3] * h * A, [Gn[0] + A * u0, -2 * Gn[0] + Gn[1] + Gn[2]]))
    Gn.append(g(Un[3]))

    Un.append(u0 + c[4] * h * pv.phiv(c[3] * h * A, [pa.zeros(pa.size(u0)), (1/2) * (- Gn[0] + Gn[1] + Gn[2] - Gn[3]), 2 * (Gn[0] - Gn[1] - Gn[2] + Gn[3])]) + c[4] * h * pv.phiv(c[4] * h * A, [Gn[0] + A * u0, -(3/2) * Gn[0] + Gn[1] + Gn[2] - (1/2) * Gn[3], Gn[0] - Gn[1] - Gn[2] + Gn[3]]))
    Gn.append(g(Un[4]))

    return u0 + h * pv.phiv(h * A, [Gn[0] + A * u0, -3 * Gn[0] - Gn[3] + 4 * Gn[4], 4 * Gn[0] + 4 * Gn[3] - 8 * Gn[4]])
